# 图像处理智能体路由配置
import os
import time
import base64
from typing import  List, Dict, Any, Union
from fastapi import APIRouter, HTTPException,  status, Form
import uuid
from pydantic import BaseModel
import json
from py_files.config import get_client,select_model
from py_files.prompt_agent_generate.get_file_type_prompt import get_agent_by_type, get_available_agent_types
import re

image_agent_router = APIRouter(prefix="/image_agent_generate", tags=["Image Agent API"])

# 图像处理请求模型
class ImageAgentRequest(BaseModel):
    agent_prompt_type: str = "通用json"
    image_file_path: str

# 图像处理响应模型
class ImageAgentResponse(BaseModel):
    id: str
    created: int
    json_result: Dict[str, Any]


# 图像处理智能体路由
@image_agent_router.post("/get_image_json", response_model=ImageAgentResponse)
async def image_agent_generate(
    agent_prompt_type: str = Form("通用json"),
    image_file_path: str = Form(...),
):
    """
    根据指定的智能体类型和图像路径生成处理结果
    
    Args:
        agent_prompt_type: 智能体类型，对应于get_file_type_prompt.py中定义的类型
        image_file_path: 图像文件路径
        
    Returns:
        处理结果的JSON响应
    """
    try:
        # if agent_prompt_type=="" or agent_prompt_type not in get_available_agent_types():
        #     agent_prompt_type = "通用json"
        prompt = get_agent_by_type("通用json")
        # 以通用json的提示词为基础提示词，其他提示词仅定义输出形式；
        if agent_prompt_type !="通用json":
            # prompt_genernal = get_agent_by_type("通用json")
            prompt_cur = get_agent_by_type(agent_prompt_type)
            prompt = prompt + prompt_cur
        if not prompt:
            available_types = get_available_agent_types()
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"未找到智能体类型 '{agent_prompt_type}'。可用类型: {', '.join(available_types)}"
            )
        
        # 验证图像文件路径是否存在
        if not os.path.exists(image_file_path):
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"图像文件路径 '{image_file_path}' 不存在"
            )
        
        # 读取图像文件
        with open(image_file_path, "rb") as image_file:
            image_data = image_file.read()
            base64_image = base64.b64encode(image_data).decode("utf-8")
        
        # 获取客户端和模型名称
        client, model_name = get_client(select_model)
        
        # 构建消息
        messages = [
            {"role": "system", "content": prompt},
            {
                "role": "user", 
                "content": [
                    {"type": "text", "text": "请分析这张图片"},
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
                ]
            }
        ]
        print("开始调用agent...")
        # 调用API
        response = client.chat.completions.create(
            model=model_name,
            messages=messages,
            temperature=0.2,
            max_tokens=1500
        )
        
        # 提取结果
        content = response.choices[0].message.content
        
        # 尝试解析JSON结果
        try:
            # 检查内容是否是JSON字符串
            if content.strip().startswith("{") and content.strip().endswith("}"):
                json_result = json.loads(content)
            else:
                # 如果不是JSON，尝试从文本中提取JSON部分
                json_matches = re.findall(r'({.*})', content, re.DOTALL)
                if json_matches:
                    json_result = json.loads(json_matches[0])
                else:
                    # 如果无法提取JSON，将结果包装在json_result中
                    json_result = {"json_result": {"raw_content": content}}
        
        except Exception as json_error:
            # 解析失败时包装原始内容
            json_result = {"json_result": {"error": str(json_error), "raw_content": content}}
        
        # 构建响应
        return ImageAgentResponse(
            id=f"imageagent-{uuid.uuid4()}",
            created=int(time.time()),
            json_result=json_result
        )
    
    except HTTPException:
        # 重新抛出HTTPException异常
        raise
    except Exception as e:
        # 处理其他所有异常
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"处理请求时发生错误: {str(e)}"
        )